import torchvision.transforms
import config
import torch
import torchvision.transforms as trans
from torch.utils.data import Dataset
import pandas as pd
from PIL import Image
import os


class YoloV1Dataset (Dataset) :
    def __init__ (
            self, data_csv, img_dir, label_dir,
            S=config.S, B=config.B, C=config.C, img_size=config.IMAGE_SIZE,
            normalize=False
    ) :
        """
        initialize the dataset, we assume the 5 to 9 depth is 0
        :param data_csv: the path to the csv file, the file struct must be [img_path,label_path]
        :param img_dir: the imgs path
        :param label_dir: the labels path
        :param S: the grid cells number per line
        :param B: the number of bboxes per grid cell
        :param C: the classes number
        :param img_size: the size of the img pass into the network
        :param normalize: whether to normalize the data, by default False
        """
        self.df = pd.read_csv (data_csv)
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.S = S
        self.B = B
        self.C = C
        self.depth = self.B * 5 + self.C
        self.img_size = img_size
        self.normalize = normalize

    def __len__ (self) :
        """
        return the length of the dataset
        :return: the length of the dataset, an integer
        """
        return len (self.df)

    def __getitem__ (self, idx) :
        """
        get the item of the dataset by idx
        :param idx: the idx of the item
        :return: return the img data and label data
        """
        # get the img path
        img_path = os.path.join (self.img_dir, self.df.iloc[idx, 0])
        label_path = os.path.join (self.label_dir, self.df.iloc[idx, 1])

        # read the img
        img = Image.open (img_path)
        label = self.read_label (label_path)

        # normalize the img
        if self.normalize :
            transforms = trans.Compose ([
                trans.Resize (self.img_size),
                trans.ToTensor (),
                trans.Normalize (
                    [0.485, 0.456, 0.406],
                    [0.22, 0.224, 0.225]
                )
            ])
            img = transforms (img)
        else :
            transforms = trans.Compose ([trans.Resize (self.img_size), trans.ToTensor ()])
            img = transforms (img)


        return img, label


    def read_label (self, label_path) :
        """
        read the label from the csv file
        :param label_path: the path to the labels
        :return: return a torch tensor with the shape of (S, S, depth) by default (7, 7, 30)
        """
        label_matrix = torch.zeros((self.S, self.S, self.depth))
        labels = []
        with open(label_path, 'r') as f :
            while True :
                line = f.readline()
                # if line is empty then break
                if not line :
                    break
                labels.append(line.strip().split())
                # [[l, x, y, w, h], [l, x, y, w, h], ...]

        for idx, label in enumerate (labels) :
            # get which grid cell is the center of the bbox in
            grid_i = int (self.S * float (label[1]))
            grid_j = int (self.S * float (label[2]))

            # get the relative x, y, w, h to the grid by percentage
            cell_x = self.S * float (label[1]) - grid_i
            cell_y = self.S * float (label[2]) - grid_j
            cell_w = self.S * float (label[3])
            cell_h = self.S * float (label[4])

            # push the data into the matrix
            label_matrix[grid_i][grid_j][0] = 1                     # the confidence
            label_matrix[grid_i][grid_j][1] = cell_x                # the x relative to grid cell
            label_matrix[grid_i][grid_j][2] = cell_y                # the y relative to grid cell
            label_matrix[grid_i][grid_j][3] = cell_w                # the width relative to grid cell
            label_matrix[grid_i][grid_j][4] = cell_h                # the height relative to grid cell
            label_matrix[grid_i][grid_j][int (label[0]) + 10] = 1   # the class score

        return label_matrix




if __name__ == '__main__' :
    # img = Image.open ("./datasets/images/000001.jpg")
    # transforms = trans.Compose ([trans.ToTensor ()])
    # print (transforms (img))
    # transforms = trans.Compose ([trans.ToTensor (), trans.Normalize([0.485, 0.456, 0.405], [0.2929, 0.224, 0.225])])
    # print (transforms (img))

    # dataset = YoloV1Dataset ("./datasets/train.csv", "./datasets/images", "./datasets/labels")
    # img, label = dataset.__getitem__(0)
    # print (img.shape)
    # print (label)

    # t = torch.zeros((3, 4, 2))
    # t[0, 1] = 2.5
    # print (t)

    train_data = YoloV1Dataset ("./datasets/train.csv", "./datasets/images", "./datasets/labels", normalize=True)
    print (train_data.__len__())
    print (train_data.__getitem__(1))